////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
//                                                                                                                            //
//                                                         WNProject                                                          //
//                                                                                                                            //
//            This file is distributed under the BSD 2-Clause open source license. See License.txt for details.               //
//                                                                                                                            //
////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////

#include "WNNetworking/inc/WNNetworkManager.h"
#include "WNContainers/inc/WNSerializer.h"
#include "WNMath/inc/WNBasic.h"

using namespace WNNetworking;
using namespace WNConcurrency;
using namespace WNMath;

WN_VOID* WNNetworkManager::IOCPThread(WN_VOID* _WNNetworkingManager) {
    WNNetworkManager* manager = reinterpret_cast<WNNetworkManager*>(_WNNetworkingManager);
    LPOVERLAPPED overLapped      = 0;
    DWORD        bytesTransfered = 0;
    ULONG_PTR    completionKey   = 0;

    while(WN_TRUE) {
        BOOL ret = GetQueuedCompletionStatus(
            manager->mIOCP,
            &bytesTransfered,
            &completionKey,
            &overLapped,
            INFINITE);
        WNConnectionData* mConnection = reinterpret_cast<WNConnectionData*>(completionKey);
        if(!ret) {
            //DOCLEANUP - ERROR OCCURRED ON CONNECTION
            continue;
        } else if (completionKey != 0 && overLapped == WN_NULL && bytesTransfered != 0) {
            if(bytesTransfered == eWNCallbackListener) {
                if(mConnection->GetCallback().Execute(mConnection->GetHandle(), mConnection->GetAddress())) {
                    mConnection->Receive();
                } else {
                    //CLOSE THE CONNECTION
                }
            }
        } else if (completionKey == 0 && overLapped == WN_NULL && bytesTransfered == 0) {
            //DOCLEANUP - SHUTDOWN REQUEST
            return(WN_NULL);
        } else if (overLapped == WN_NULL && bytesTransfered == 0) {
            //DOCLEANUP - Connection was closed
            continue;
        } else {
            switch(mConnection->GetOperation(overLapped)){
                case WNConnectionData::WriteOp:
                    mConnection->ProcessWrite();
                    break;
                case WNConnectionData::ReadOp:
                    mConnection->ProcessRead(bytesTransfered);
                    mConnection->Receive();
                    break;
                case WNConnectionData::NoOp:
                    WN_RELEASE_ASSERT_DESC(WN_FALSE, "Bad operation on this socket");
                    break;
            }
            //1) Figure out what operation just finished
            //2) Process next operation
            //Since we want to be able to process multiple reads/writes
            //This may change significantly with UDP, since we cannot control WHO the data is coming from
        }
    }
}

WN_VOID* WNNetworkManager::ListenThread(WN_VOID* _WNNetworkingManager) {
    WNNetworkManager* manager = reinterpret_cast<WNNetworkManager*>(_WNNetworkingManager);
    WSANETWORKEVENTS events;
    while(WAIT_OBJECT_0 != WaitForSingleObject(manager->mShutdownEvent, 0))
    {
        if (WSA_WAIT_TIMEOUT != WSAWaitForMultipleEvents(1, &manager->mAcceptEvent, FALSE, 1, FALSE))
        {
            manager->mListenSection.Lock();
            for(std::map<SOCKET, WNCallback2<WN_BOOL, WNConnectionHandle, const WN_CHAR*> >::iterator i = manager->mListeningSockets.begin(); i != manager->mListeningSockets.end(); ++i)
            {
                WSAEnumNetworkEvents(i->first, manager->mAcceptEvent, &events);
                if(events.lNetworkEvents & FD_ACCEPT && (0 == events.iErrorCode[FD_ACCEPT_BIT])) {
                    SOCKET s = i->first;
                    sockaddr_in address;
                    int sockAddrSize = sizeof(address);
                    //Accept remote connection attempt from the client
                    SOCKET socket = accept(i->first, (sockaddr*)&address, &sockAddrSize);
                    if(INVALID_SOCKET == socket) {
                        //HANLDE BAD SOCKET CONNECTION
                        continue;
                    }
                    
                    WNConnectionHandle handleNum = 0;
                    while((handleNum = WNAtomicCompareSwap(&manager->mCurrentHandleVal, manager->mCurrentHandleVal + 1, manager->mCurrentHandleVal)) != manager->mCurrentHandleVal - 1);
                    WN_SIZE_T handleType = eWNSocketFromClient;
                    handleType = handleType << CONNECTION_SHIFT;
                    handleNum |= handleType;  

                    WNConnectionData* data = WN_NEW(WNConnectionData(*manager));
                    data->SetSocket(socket);
                    data->SetAddress(inet_ntoa(address.sin_addr));

                    HANDLE temp = CreateIoCompletionPort((HANDLE)socket, manager->mIOCP, (ULONG_PTR)data, 0);

                    if(temp == NULL) {
                        WN_DELETE(data);
                        closesocket(socket);
                        continue;
                    } else {
                        manager->mConnectionLock.Lock();
                        manager->mHandles[handleNum] = data;
                        manager->mConnectionLock.Unlock();
                        data->SetHandle(handleNum);
                        data->SetCallback(i->second);
                        PostQueuedCompletionStatus(manager->mIOCP, eWNCallbackListener, reinterpret_cast<ULONG_PTR>(data), WN_NULL);
                    }

                } else if(events.lNetworkEvents & FD_ACCEPT && (0 != events.iErrorCode[FD_ACCEPT_BIT])) {
                    //HANDLE ERROR (on socket s)
                    SOCKET s = i->first;
                    manager->mListeningSockets[i->first] = manager->mListeningSockets[manager->mListeningSockets.size() - 1];
                    manager->mListeningSockets.erase(i->first);
                    manager->mListenSection.Unlock();
                    manager->mListenSection.Lock();
                } else {
                }
            }
            
            manager->mListenSection.Unlock();
        }
        //HANDLE any pushed operations here (things that are queued up.. might not even be any)
    }
    return(WN_NULL);
}


WNNetworkManager::WNNetworkManager():
    mMaxThreads(0),
    mIOCP(INVALID_HANDLE_VALUE),
    mAcceptEvent(INVALID_HANDLE_VALUE),
    mInitializationState(WN_NOT_STARTED),
    mCurrentHandleVal(1){
}

WNNetworkingError WNNetworkManager::Initialize(WN_UINT32 _numWorkerThreads) {
    mMaxThreads = _numWorkerThreads;
    int wsaRet = WSAStartup(MAKEWORD(2, 2), &mWSAData);
    if(wsaRet != NO_ERROR) {
        Cleanup();
        return(WN_NETWORKING_MANAGER_SYSTEM_INITIALIZATION_FAILED);
    }
    mInitializationState = WN_WSA_STARTUP;
    mAcceptEvent = WSACreateEvent();
    if (WSA_INVALID_EVENT == mAcceptEvent)
    {
        Cleanup();
        return(WN_NETWORKING_MANAGER_SYSTEM_INITIALIZATION_FAILED);
    }
    mInitializationState = WN_ACCEPT_EVENT_CREATED;
    mShutdownEvent = CreateEvent(NULL, TRUE, FALSE, NULL);
    if(mShutdownEvent == INVALID_HANDLE_VALUE) {
        Cleanup();
        return(WN_NETWORKING_MANAGER_SYSTEM_INITIALIZATION_FAILED);
    }
    mInitializationState = WN_SHUTDOWN_EVENT_CREATED;

    mIOCP = CreateIoCompletionPort(INVALID_HANDLE_VALUE, WN_NULL, 0, 0);
    if(mIOCP == INVALID_HANDLE_VALUE) {
        Cleanup();
        return(WN_NETWORKING_MANAGER_SYSTEM_INITIALIZATION_FAILED);
    }

    mInitializationState = WN_IOCP_CREATED;

    for(WN_UINT32 i = 0; i < _numWorkerThreads; ++i) {
        WNThread* thread = new WNThread();

		// need to add stack size stuff for 1 MB stack
        if(!thread->SetStartRoutine(&WNNetworkManager::IOCPThread, this)) {
            Cleanup();

            return(WN_NETWORKING_MANAGER_THREAD_CREATION_ERROR);
        } else {
			thread->Execute();
            mThreads.push_back(thread);
        }
    }

    mInitializationState = WN_THREADS_CREATED;

		// need to add stack size stuff for 1 MB stack
    if(!mListenThread.SetStartRoutine(&WNNetworkManager::ListenThread, this)) {
        Cleanup();

        return(WN_NETWORKING_MANAGER_THREAD_CREATION_ERROR);
    } else {
		mListenThread.Execute();

		mInitializationState = WN_LISTEN_THREAD_CREATED;
	}

    mInitializationState = WN_INITIALIZATION_COMPLETE;

    return(WN_OK);
}

WN_VOID WNNetworkManager::Cleanup() {
    switch((mInitializationState)) {
        case WN_INITIALIZATION_COMPLETE:
        case WN_LISTEN_THREAD_CREATED:
            //Ask all threads to start shutting down
            SetEvent(mShutdownEvent);

            mListenThread.Join();
        case WN_THREADS_CREATED:
            for(WN_SIZE_T i = 0; i < mThreads.size(); ++i) {
                PostQueuedCompletionStatus(mIOCP, 0, NULL, 0);
            }

			for (WN_UINT32 i = 0; i < static_cast<WN_UINT32>(mThreads.size()); ++i) {
				mThreads[i]->Join();

				WN_DELETE(mThreads[i]);
			}
        case WN_IOCP_CREATED:
            CloseHandle(mIOCP);
        case WN_SHUTDOWN_EVENT_CREATED:
            CloseHandle(mShutdownEvent);
        case WN_ACCEPT_EVENT_CREATED:
            WSACloseEvent(mAcceptEvent);
        case WN_WSA_STARTUP:
            WSACleanup();
        case WN_NOT_STARTED:
            break;
    }
    for(std::map<WNConnectionHandle, WNConnectionData*>::iterator i = mHandles.begin(); i != mHandles.end(); ++i){
        WN_DELETE(i->second);
    }
    for(std::map<WNConnectionHandle, Group*>::iterator i = mGroups.begin(); i != mGroups.end(); ++i){
        WN_DELETE(i->second);
    }
    mHandles.clear();
    mGroups.clear();
}

WNNetworkingError WNNetworkManager::ConnectTo(WNConnectionHandle& _outHandle, WNConnectionType _type, const WN_CHAR* _target, WN_SIZE_T _port) {
    WN_RELEASE_ASSERT_DESC(_type == eWNReliable, "WNNetworking does not support unreliable connections... YET");
    WN_RELEASE_ASSERT_DESC(_port <= 0xFFFF, "Error port must be within range");
    SOCKET ClientSocket = WSASocket(AF_INET, SOCK_STREAM, 0, NULL, 0, WSA_FLAG_OVERLAPPED);
    if (INVALID_SOCKET == ClientSocket) {
        return(WN_NETWORKING_MANAGER_CANNOT_CREATE_SOCKET);
    }
    struct sockaddr_in address;
    struct hostent* server;

    server = gethostbyname(_target);
    if(server == NULL) {
        closesocket(ClientSocket);
        return(WN_NETWORKING_MANAGER_BAD_HOSTNAME);
    }
    ZeroMemory((char*) &address, sizeof(sockaddr_in));
    address.sin_family = AF_INET;
    CopyMemory((char*) &address.sin_addr.s_addr, (char*)server->h_addr, server->h_length);
    address.sin_port = htons(static_cast<u_short>(_port));

    WNConnectionData* dat = WN_NEW(WNConnectionData(*this));
    dat->SetSocket(ClientSocket);
    dat->SetAddress(_target);
    WN_SIZE_T handleNum = 0;
    while((handleNum = WNAtomicCompareSwap(&mCurrentHandleVal, mCurrentHandleVal + 1, mCurrentHandleVal)) != mCurrentHandleVal - 1);
    WN_SIZE_T handleType = eWNSocketClient;
    handleType = handleType << CONNECTION_SHIFT;
    handleNum |= handleType;
  

    if (SOCKET_ERROR == connect(ClientSocket, reinterpret_cast<const struct sockaddr*>(&address), sizeof(address))){
        WN_DELETE(dat);
        closesocket(ClientSocket);
        return(WN_NETWORKING_MANAGER_CANNOT_CONNECT);
    }

    HANDLE temp = CreateIoCompletionPort((HANDLE)ClientSocket, mIOCP, (ULONG_PTR)dat, 0);
    if(temp == NULL) {
        WN_DELETE(dat);
        closesocket(ClientSocket);
        return(WN_NETWORKING_MANAGER_CANNOT_ASSOCIATE);
    }

    mConnectionLock.Lock();
    mHandles[handleNum] = dat;
    mConnectionLock.Unlock();
    dat->Receive();
    _outHandle = handleNum;
    return(WN_OK);
}
WNNetworkingError WNNetworkManager::CreateListener(WNConnectionHandle& _outHandle, WNConnectionType _type, WN_UINT16 _port, WNCallback2<WN_BOOL, WNConnectionHandle, const WN_CHAR*> _callback) {
    WN_RELEASE_ASSERT_DESC(_type == eWNReliable, "WNNetworking does not support unreliable connections... YET");
    SOCKET listenSocket = WSASocket(AF_INET, SOCK_STREAM, 0, NULL, 0, WSA_FLAG_OVERLAPPED);
    if(INVALID_SOCKET == listenSocket) {
        return(WN_NETWORKING_MANAGER_CANNOT_CREATE_SOCKET);
    }

    sockaddr_in ServerAddress;
    ZeroMemory(&ServerAddress, sizeof(sockaddr_in));
    ServerAddress.sin_family = AF_INET;
    ServerAddress.sin_addr.s_addr = INADDR_ANY;
    ServerAddress.sin_port = htons(_port);

    if(SOCKET_ERROR == bind(listenSocket, (struct sockaddr *) &ServerAddress, sizeof(ServerAddress))) {
        closesocket(listenSocket);
        return(WN_NETWORKING_MANAGER_CANNOT_CREATE_SOCKET);
    }

    if(SOCKET_ERROR == listen(listenSocket, SOMAXCONN)) {
        closesocket(listenSocket);
        return(WN_NETWORKING_MANAGER_CANNOT_CREATE_SOCKET);
    }

    

    WNConnectionData* dat = WN_NEW(WNConnectionData(*this));
    dat->SetSocket(listenSocket);

    WN_SIZE_T handleNum = 0;
    while((handleNum = WNAtomicCompareSwap(&mCurrentHandleVal, mCurrentHandleVal + 1, mCurrentHandleVal)) != mCurrentHandleVal - 1);
    WN_SIZE_T handleType = eWNSocketServer;
    handleType = handleType << CONNECTION_SHIFT;
    handleNum |= handleType;
    mConnectionLock.Lock();
    mHandles[handleNum] = dat;
    mConnectionLock.Unlock();

    mListenSection.Lock();
    mListeningSockets[listenSocket] = _callback;
    mListenSection.Unlock();

    if(SOCKET_ERROR == WSAEventSelect(listenSocket, mAcceptEvent, FD_ACCEPT)) {
        mConnectionLock.Lock();
        mHandles.erase(handleNum);
        mHandles[handleNum] = dat;
        mListenSection.Lock();
        mListeningSockets.erase(listenSocket);
        mListenSection.Unlock();
        WN_DELETE(dat);
        mConnectionLock.Unlock();
        closesocket(listenSocket);
        return(WN_NETWORKING_MANAGER_CANNOT_LISTEN);
    }

    _outHandle = handleNum;
    return(WN_OK);
}

WN_VOID WNNetworkManager::WNConnectionData::ProcessRead(WN_SIZE_T _transferred) {
    WN_SIZE_T processedBytes = 0;
    WNContainers::WNResourcePointer<WNBufferResource> p = mReadBuffers[0];
    mReadHead += _transferred;
    while(processedBytes != _transferred)
    {
        WN_DEBUG_ASSERT(mReadHead <= WNContainers::WNDataBuffer::mMaxWNDataBufferWrite);
        if(mExpectedData == 0) {
            if(mReadHead - mProcessedHead + mLeftoverData >= 8) {
                WN_CHAR dataPlace[8];
                for(WN_SIZE_T i = 0; i < mLeftoverData; ++i) {
                    dataPlace[i] = mLeftover[i];
                }
                for(WN_SIZE_T i = 0; i < (8 - mLeftoverData); ++i){
                    dataPlace[mLeftoverData + i] = *(p->GetBaseLocation() + mProcessedHead + i);
                }
                
                WN_UINT32 num = *reinterpret_cast<WN_UINT32*>(dataPlace);
                WN_UINT32 callback = *reinterpret_cast<WN_UINT32*>(dataPlace + 4);
                FromBigEndian(callback);
                FromBigEndian(num);
				mExpectedData = num;
                mCallbackFunction = callback;
                mProcessedHead += (8 - mLeftoverData);
                mProcessedData += (8 - mLeftoverData);
                processedBytes += (8 - mLeftoverData);
                mLeftoverData = 0;
            } else {
                if (mReadHead - mProcessedHead + mLeftoverData > 8) {
                    WN_DEBUG_ASSERT(mProcessedHead == 0);
                    WN_CHAR* leftoverData = mCurrentReadBuffer.GetLastBuffer();
                    WN_CHAR dataPlace[8];
                    WN_SIZE_T ctr = 0;
                    for(ctr = 0; ctr < mLeftoverData; ctr++){
                        dataPlace[ctr] = leftoverData[WNContainers::WNDataBuffer::mMaxWNDataBufferWrite - mLeftoverData + ctr];
                    }
                    for(WN_SIZE_T i = 0; i < (8 - ctr); ++i){
                        dataPlace[(8 - ctr) + i] = *(p->GetBaseLocation() + mProcessedHead + i);
                    }
                    WN_UINT32 num = *reinterpret_cast<WN_UINT32*>(dataPlace);
                    FromBigEndian(num);
                    mExpectedData = num;
                    mProcessedHead += 8 - mLeftoverData;
                    mProcessedData += 8 - mLeftoverData;
                    processedBytes += 8 - mLeftoverData;
                    mLeftoverData = 0;
                } else {
                    //We have some left-over bytes
                    WN_SIZE_T leftover = mReadHead - mProcessedHead;
                    for(WN_SIZE_T i = 0; i < leftover; ++i) {
                        mLeftover[mLeftoverData + i] = *(p->GetBaseLocation() + mProcessedHead + i);
                    }
                    mLeftoverData += leftover;
                    mProcessedHead = mReadHead;
                    mProcessedData += leftover;
                    break;
                }
            }
        }
        if(mExpectedData != 0) {
            WN_SIZE_T maxProcess = WNMin(mReadHead - mProcessedHead, mExpectedData - mProcessedData);
            mProcessedData += maxProcess; //process as much as possible;
            if(mProcessedData == mExpectedData) {
                mCurrentReadBuffer.AppendBuffer(p, maxProcess, mProcessedHead);
                WNCallback2<WN_VOID, WNConnectionHandle, WNNetworkReadBuffer&>* cb =  mManager.GetCallback(mCallbackFunction);
                WN_RELEASE_ASSERT_DESC(cb != WN_NULL, "Invalid callback for buffer");
                cb->Execute(mHandle, mCurrentReadBuffer);
                mCurrentReadBuffer.Clear();
                mProcessedData = 0;
                mExpectedData = 0;
                mCallbackFunction = 0;
            }
            
            mProcessedHead += maxProcess;
            processedBytes += maxProcess;
        }
    }

    if(mReadHead == WNContainers::WNDataBuffer::mMaxWNDataBufferWrite) {
        mReadBuffers[0] = WNContainers::WNAllocateResource<WNBufferResource>(mManager);
        mProcessedHead = mReadHead = 0;
    } else {
        mReadBuffers[0]->AddData(_transferred);
    }
}

WNNetworkingError WNNetworkManager::SendBuffer(WNConnectionHandle _handle, WNNetworkWriteBuffer& buffer) {
    WN_SIZE_T buffInfo = (_handle >> CONNECTION_SHIFT) & CONNECTION_MASK;
    buffer.FlushWrite();
    switch(buffInfo) {
        case eWNSocketServer:
            //HANDLE SENDING TO ALL THE PEOPLE CONNECTED TO THE LISTENING SOCKET
            break;
        case eWNSocketFromClient: //These are handled the same way
        case eWNSocketClient:{
            WNConnectionData* handle = mHandles[_handle];
            WN_DEBUG_ASSERT_DESC(handle, "Handle does not exist that you are trying to send to");
            handle->AppendSendBuff(buffer);
            break;
        }
        case eWNSocketGroup:
            Group* g = mGroups[_handle];
            for(std::vector<WNConnectionHandle>::iterator i = g->mGroupedHandles.begin(); i != g->mGroupedHandles.end(); ++i){
                SendBuffer(*i, buffer);
            }
            break;
    }
    return(WN_OK);
}

WN_BOOL WNNetworkManager::WNConnectionData::Receive() {
    DWORD bytes, flags;
    bytes = flags = 0;
    mReadBuffers[0]->PrepareRead();
    int received = WSARecv(mSocket, mReadBuffers[0]->GetReadWinBuf(), 1, &bytes, &flags, &mReceiveOverlap, NULL);
    int lastError = WSAGetLastError();

    if(SOCKET_ERROR == received && (WSA_IO_PENDING != lastError)) {
        return(WN_FALSE);
    }
    return(WN_TRUE);
}

WN_VOID WNNetworkManager::WNConnectionData::AppendSendBuff(WNNetworkWriteBuffer& _buff) {
    mWriteBufferLock.Lock();
    mWriteBuffers.push_back(_buff);
    if(mWriteBuffers.size() == 1){
        Send();
    }
    mWriteBufferLock.Unlock();
}

WN_VOID WNNetworkManager::WNConnectionData::Send(){
    if(mSending){
        return;
    }
    mSending = WN_TRUE;
    WNNetworkWriteBuffer& buff = mWriteBuffers.front();
    const std::deque<WNContainers::WNResourcePointer<WNNetworking::WNBufferResource> >& sendBuffs = buff.GetChunks();
    mWSAWriteBuffers.clear();
    for(std::deque<WNContainers::WNResourcePointer<WNNetworking::WNBufferResource> >::const_iterator i = sendBuffs.begin(); i != sendBuffs.end(); ++i){
        mWSAWriteBuffers.push_back(*(*i)->GetWriteWinBuf());
    }

    DWORD bytes;
    int sent = WSASend(mSocket, &mWSAWriteBuffers[0], static_cast<DWORD>(sendBuffs.size()), &bytes, 0, &mSendOverlap, NULL);
    if(SOCKET_ERROR == sent && WSA_IO_PENDING != WSAGetLastError()) {
        //Handle the fact that we are sending data to someone we shouldn't
    }
}

WN_VOID WNNetworkManager::WNConnectionData::ProcessWrite() {
    mSending = WN_FALSE;
    mWriteBufferLock.Lock();
    mWriteBuffers.pop_front();
    if(mWriteBuffers.size() > 0){
        Send();
    }
    mWriteBufferLock.Unlock();
}

WN_VOID WNNetworkManager::InitializeBuffer(WNNetworkWriteBuffer& _buffer, WN_UINT32 _number) {
    WN_UINT32 num = 0;
    _buffer.Serialize(0, WNContainers::WNSerializer<WN_UINT32>(num));
    _buffer.Serialize(0, WNContainers::WNSerializer<WN_UINT32>(_number));
}

WNCallback2<WN_VOID, WNConnectionHandle, WNNetworkReadBuffer&>* WNNetworkManager::GetCallback(WN_UINT32 _identifier) {
    if(mCallbacks.find(_identifier) != mCallbacks.end()) {
        return(&mCallbacks[_identifier]);
    }
    return(WN_NULL);
}

WN_VOID WNNetworkManager::SetCallback(WN_UINT32 _identifier, WNCallback2<WN_VOID, WNConnectionHandle, WNNetworkReadBuffer&> _buffer) {
    WN_DEBUG_ASSERT(mCallbacks.find(_identifier) == mCallbacks.end());
    mCallbacks[_identifier] = _buffer;
}